package crypto

import (
	"bytes"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"encoding/base64"
	"errors"
	"fmt"
	"math/big"
	"net/url"
)

type RSA struct {
	privateKey *rsa.PrivateKey
	publicKey  *rsa.PublicKey
}

// NewRSACrypto 返回rsa结构体实例
func NewRSACrypto(pubKey *rsa.PublicKey, prvKey *rsa.PrivateKey) *RSA {
	return &RSA{
		privateKey: prvKey,
		publicKey:  pubKey,
	}
}

// EncryptWithPrivateKey 私钥加密
func (r *RSA) EncryptWithPrivateKey(data []byte) ([]byte, error) {
	keySize, srcSize := r.privateKey.Size(), len(data)
	//logs.Debug("密钥长度：", keySize, "\t密文长度：\t", srcSize)
	offSet, once := 0, keySize-11
	var buffer = bytes.Buffer{}
	for offSet < srcSize {
		endIndex := offSet + once
		if endIndex > srcSize {
			endIndex = srcSize
		}
		bytesOnce, err := rsa.SignPKCS1v15(rand.Reader, r.privateKey, crypto.Hash(0), data[offSet:endIndex])
		if err != nil {
			return nil, err
		}
		buffer.Write(bytesOnce)
		offSet = endIndex
	}
	return buffer.Bytes(), nil
}

// DecryptWithPublicKey 公钥解密
func (r *RSA) DecryptWithPublicKey(data []byte) ([]byte, error) {
	keySize, srcSize := r.publicKey.Size(), len(data)
	var offSet = 0
	var buffer = bytes.Buffer{}
	for offSet < srcSize {
		endIndex := offSet + keySize
		if endIndex > srcSize {
			endIndex = srcSize
		}
		decData, err := pubKeyDecrypt(r.publicKey, crypto.Hash(0), nil, data[offSet:endIndex])
		if err != nil {
			return nil, err
		}
		buffer.Write(decData)
		offSet = endIndex
	}
	return buffer.Bytes(), nil
}

// copy&modified from crypt/rsa/pkcs1v5.go
func pubKeyDecrypt(pub *rsa.PublicKey, hash crypto.Hash, hashed []byte, sig []byte) (out []byte, err error) {
	hashLen, prefix, err := pkcs1v15HashInfo(hash, len(hashed))
	if err != nil {
		return nil, err
	}

	tLen := len(prefix) + hashLen
	k := (pub.N.BitLen() + 7) / 8
	if k < tLen+11 {
		return nil, fmt.Errorf("length illegal")
	}

	c := new(big.Int).SetBytes(sig)
	m := encrypt(new(big.Int), pub, c)
	em := leftPad(m.Bytes(), k)
	out = unLeftPad(em)

	err = nil
	return
}

// copy from crypt/rsa/pkcs1v5.go
func encrypt(c *big.Int, pub *rsa.PublicKey, m *big.Int) *big.Int {
	e := big.NewInt(int64(pub.E))
	c.Exp(m, e, pub.N)
	return c
}

// copy from crypt/rsa/pkcs1v5.go
var hashPrefixes = map[crypto.Hash][]byte{
	crypto.MD5:       {0x30, 0x20, 0x30, 0x0c, 0x06, 0x08, 0x2a, 0x86, 0x48, 0x86, 0xf7, 0x0d, 0x02, 0x05, 0x05, 0x00, 0x04, 0x10},
	crypto.SHA1:      {0x30, 0x21, 0x30, 0x09, 0x06, 0x05, 0x2b, 0x0e, 0x03, 0x02, 0x1a, 0x05, 0x00, 0x04, 0x14},
	crypto.SHA224:    {0x30, 0x2d, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x04, 0x05, 0x00, 0x04, 0x1c},
	crypto.SHA256:    {0x30, 0x31, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x01, 0x05, 0x00, 0x04, 0x20},
	crypto.SHA384:    {0x30, 0x41, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x02, 0x05, 0x00, 0x04, 0x30},
	crypto.SHA512:    {0x30, 0x51, 0x30, 0x0d, 0x06, 0x09, 0x60, 0x86, 0x48, 0x01, 0x65, 0x03, 0x04, 0x02, 0x03, 0x05, 0x00, 0x04, 0x40},
	crypto.MD5SHA1:   {}, // A special TLS case which doesn't use an ASN1 prefix.
	crypto.RIPEMD160: {0x30, 0x20, 0x30, 0x08, 0x06, 0x06, 0x28, 0xcf, 0x06, 0x03, 0x00, 0x31, 0x04, 0x14},
}

// copy from crypt/rsa/pkcs1v5.go
func pkcs1v15HashInfo(hash crypto.Hash, inLen int) (hashLen int, prefix []byte, err error) {
	// Special case: crypto.Hash(0) is used to indicate that the data is
	// signed directly.
	if hash == 0 {
		return inLen, nil, nil
	}

	hashLen = hash.Size()
	if inLen != hashLen {
		return 0, nil, errors.New("crypto/rsa: input must be hashed message")
	}
	prefix, ok := hashPrefixes[hash]
	if !ok {
		return 0, nil, errors.New("crypto/rsa: unsupported hash function")
	}
	return
}

// copy from crypt/rsa/pkcs1v5.go
func leftPad(input []byte, size int) (out []byte) {
	n := len(input)
	if n > size {
		n = size
	}
	out = make([]byte, size)
	copy(out[len(out)-n:], input)
	return
}

func unLeftPad(input []byte) (out []byte) {
	n := len(input)
	t := 2
	for i := 2; i < n; i++ {
		if input[i] == 0xff {
			t = t + 1
		} else {
			if input[i] == input[0] {
				t = t + int(input[1])
			}
			break
		}
	}
	out = make([]byte, n-t)
	copy(out, input[t:])
	return
}

// EncryptWithPublicKey 公钥加密
// pkcs 1# v1.5
func (r *RSA) EncryptWithPublicKey(data []byte) ([]byte, error) {
	keySize, dataSize := r.publicKey.Size(), len(data)
	offset, once := 0, keySize-11
	buffer := bytes.Buffer{}
	for offset < dataSize {
		endIndex := offset + once
		if endIndex > dataSize {
			endIndex = dataSize
		}
		// 加密一部分
		bytesOnce, err := rsa.EncryptPKCS1v15(rand.Reader, r.publicKey, data[offset:endIndex])
		if err != nil {
			return nil, err
		}
		buffer.Write(bytesOnce)
		offset = endIndex
	}
	return buffer.Bytes(), nil
}

// DecryptWithPrivateKey 私钥解密
// 分段
// pkcs 1# v1.5
func (r *RSA) DecryptWithPrivateKey(data []byte) ([]byte, error) {
	keySize, srcSize := r.privateKey.Size(), len(data)
	var offSet = 0
	var buffer = bytes.Buffer{}
	for offSet < srcSize {
		endIndex := offSet + keySize
		if endIndex > srcSize {
			endIndex = srcSize
		}
		bytesOnce, err := rsa.DecryptPKCS1v15(rand.Reader, r.privateKey, data[offSet:endIndex])
		if err != nil {
			return nil, err
		}
		buffer.Write(bytesOnce)
		offSet = endIndex
	}
	return buffer.Bytes(), nil
}

// EncryptWithPrvKOutputBase64 私钥加密
// 输入base64
func (r *RSA) EncryptWithPrvKOutputBase64(data []byte) (string, error) {
	content := url.QueryEscape(string(data))
	en, err := r.EncryptWithPrivateKey([]byte(content))
	if err != nil {
		return "", err
	}
	return base64.URLEncoding.EncodeToString(en), nil
}

// DecryptWithPubKeyFromBase64 公钥解密
// 输入base64
func (r *RSA) DecryptWithPubKeyFromBase64(data string) ([]byte, error) {
	dataByte, err := base64.URLEncoding.DecodeString(data)
	if err != nil {
		return nil, err
	}
	dataByte, err = r.DecryptWithPublicKey(dataByte)
	if err != nil {
		return nil, err
	}
	result, err := url.QueryUnescape(string(dataByte))
	return []byte(result), err
}

// EncryptWithPubKeyOutputBase64 公钥加密
// 输出base64
func (r *RSA) EncryptWithPubKeyOutputBase64(data []byte) (string, error) {
	content := url.QueryEscape(string(data))
	en, err := r.EncryptWithPublicKey([]byte(content))
	if err != nil {
		return "", err
	}
	return base64.URLEncoding.EncodeToString(en), nil
}

// DecryptWithPrvKeyFromBase64 私钥解密
// 输入base64
func (r *RSA) DecryptWithPrvKeyFromBase64(data string) ([]byte, error) {
	dataByte, err := base64.URLEncoding.DecodeString(data)
	if err != nil {
		return nil, err
	}
	dataByte, err = r.DecryptWithPrivateKey(dataByte)
	if err != nil {
		return nil, err
	}
	result, err := url.QueryUnescape(string(dataByte))
	return []byte(result), err
}
